import gradio as gr
import torch
from torchvision import datasets

model = torch.load('fashion.pth')
data_folder = 'E:/study_code/torch_study/aigc/data'
fmnist = datasets.FashionMNIST(data_folder,download=True,train=True)
# cls = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
def image_classifier(inp):
    model.eval()
    tr = torch.from_numpy(inp)
    tr = tr.float() / 255
    tr = tr.view(-1, 1, 28, 28)
    tr = torch.tensor(tr)
    result = model(tr.to('cuda'))
    max_values, argmaxes = result.max(-1)
    vl = str(fmnist.classes[argmaxes[0].cpu().detach().numpy()])
    return vl

demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
demo.launch()
